import math
from QKDNetwork import QKDNetwork
from docplex.mp.model import Model

class LP:
    def __init__(self, net: QKDNetwork, find_ultimate_time_slot = False):
        self.net = net
        self.node_num = net.num_nodes + 1 # 用户方便使用ID索引节点
        self.sd_num = len(net.sd_list)
        self.max_candidate_size = net.max_candidate_size
        self.estimate_max_time_slot = net.estimate_max_time_slot
        self.selected_sole_paths = {}
        self.have_designed_sole_path = False
        # 是否要对于选定的路径再次进行Cplex搜索完成时隙
        self.find_ultimate_time_slot = find_ultimate_time_slot
        self.find_min_feasible_T()
    
    def find_min_feasible_T(self):
        low, high = 0, self.estimate_max_time_slot - 1
        self.best_feasible_T = None # 用于存储最佳的Time Slot
        self.best_model = None  # 用于存储最佳模型

        while low <= high:
            mid = (low + high) // 2
            model = self.build_model(mid)
            sol = model.solve()
            if sol is not None:  # 如果找到可行解
                # print(f"Feasible solution found at T = {mid}")
                self.best_feasible_T = mid
                self.best_model = model  # 保存当前模型的拷贝
                high = mid - 1  # 尝试更小的T
            else:
                # print(f"No feasible solution found at T = {mid}")
                low = mid + 1  # 增大T值尝试

        if self.best_feasible_T is not None and self.best_model is not None:
            # print(f"The minimum feasible T is {self.best_feasible_T}")
            # 打印 z 变量的值
            for i in range(self.sd_num):
                for j in range(self.max_candidate_size):
                    for t in range(self.best_feasible_T + 1):
                        var_name = f'z_{i}_{j}_{t}'
                        var = self.best_model.get_var_by_name(var_name)
                        if var is not None and var.solution_value != 0 and t == self.best_feasible_T - 1:
                            if(self.have_designed_sole_path):
                                pass
                                # print(f"{var_name} = {var.solution_value}", end=", ")
            if(self.have_designed_sole_path == False):
                self.determineSolePath()
        else:
            print("No feasible T found within the estimated range")

    def determineSolePath(self):
        # 确保最佳模型已解决并可用
        assert self.best_model is not None
        # 从解决方案中提取在时间T-1（最后可能的时间槽）的 z 变量
        self.selected_sole_paths = {}
        for i in range(self.sd_num):
            for j in range(self.max_candidate_size):
                var_name = f'z_{i}_{j}_{self.best_feasible_T - 1}'
                var = self.best_model.get_var_by_name(var_name)
                if var.solution_value > 0:
                    if i in self.selected_sole_paths:
                        # 如果已经选择了一条路径，则选择具有最大z值的那条
                        current_j, current_value = self.selected_sole_paths[i]
                        if var.solution_value > current_value:
                            self.selected_sole_paths[i] = (j, var.solution_value)
                    else:
                        self.selected_sole_paths[i] = (j, var.solution_value)
        # print(f"全部路径")
        # print(self.net.candidate_paths)
        # print(f"选中的路径")
        # print(self.selected_sole_paths)
        if(self.find_ultimate_time_slot):
            self.have_designed_sole_path = True
            self.find_min_feasible_T()

    def build_model(self, fixed_T: int):
        determined_sole_path = self.selected_sole_paths
        # 构建模型
        model = Model(name="qkd_manet_heterogeneous")
        x_dim1, x_dim2, x_dim3, x_dim4, x_dim5 = self.sd_num, self.max_candidate_size, self.node_num, self.node_num, fixed_T + 1
        x = model.continuous_var_dict(
            [(i, j, k, l, m)    for i in range(x_dim1)
                                for j in range(x_dim2)
                                for k in range(x_dim3)
                                for l in range(x_dim4)
                                for m in range(x_dim5)], lb=0, ub=1, name='x'
        )

        y_dim1, y_dim2, y_dim3, y_dim4, y_dim5 = self.sd_num, self.max_candidate_size, self.node_num, self.node_num, fixed_T + 1
        y = model.continuous_var_dict(
            [(i, j, k, l, m)    for i in range(y_dim1)
                                for j in range(y_dim2)
                                for k in range(y_dim3)
                                for l in range(y_dim4)
                                for m in range(y_dim5)], lb=0, ub=1, name='y'
        )

        z_dim1, z_dim2, z_dim3 = self.sd_num, self.max_candidate_size, fixed_T + 1
        z = model.continuous_var_dict(
            [(i, j, k)  for i in range(z_dim1)
                        for j in range(z_dim2)
                        for k in range(z_dim3)], lb=0, ub=1, name='z'
        )

        w_dim1, w_dim2 = self.sd_num, fixed_T + 1
        w = model.continuous_var_dict(
            [(i, j)  for i in range(w_dim1)
                        for j in range(w_dim2)], lb=0, ub=1, name='w'
        )

        # 2a
        for u in self.net.G.nodes:
            for t in range(fixed_T):
                x_ijuvt = model.linear_expr()
                for i in range(self.sd_num):
                    for j in range(len(self.net.candidate_paths[i]["paths"])):
                        for v in list(self.net.G.neighbors(u)):
                            x_ijuvt += x[i, j, u, v, t]
                x_ijuvt -= self.net.G.nodes[u]["transmitter"]
                model.add_constraint(x_ijuvt <= 0, ctname="x_less_su")
        
        # 2b
        for i in range(self.sd_num):
            for j in range(len(self.net.candidate_paths[i]["paths"])):
                for v in self.net.G.nodes:
                    # 使这个v点，以及下面的u点，属于预置路径
                    for t in range(fixed_T):
                        for u in list(self.net.G.neighbors(v)):
                            y_ijuvt = model.linear_expr()
                            # 将 u,v Swap
                            y_ijuvt += y[i, j, u, v, t]
                            for k in range(t + 1):
                                if(self.net.G.nodes[u]["transmitter_rate"] != 0):
                                    y_ijuvt -= x[i, j, u, v, k] / math.ceil(self.net.sd_list[i][2] / self.net.G.nodes[u]["transmitter_rate"])
                                if(self.net.G.nodes[v]["transmitter_rate"] != 0):
                                    y_ijuvt -= x[i, j, v, u, k] / math.ceil(self.net.sd_list[i][2] / self.net.G.nodes[v]["transmitter_rate"])
                            model.add_constraint(y_ijuvt <= 0, ctname="y_less_x")
        
        # 更新模型约束来强制选择单一路径
        for i, (selected_j, value) in determined_sole_path.items():
            for j in range(self.max_candidate_size):
                if j != selected_j:
                    for t in range(fixed_T):
                        model.add_constraint(z[i, j, t] == 0)

        # 2c
        for i in range(self.sd_num):
            for j in range(len(self.net.candidate_paths[i]["paths"])):
                for t in range(fixed_T):
                    z_ijt = model.linear_expr()
                    z_ijt += z[i, j, t]
                    candidate_path = self.net.candidate_paths[i]["paths"][j]
                    candidate_path_len = len(candidate_path)
                    for path_start_index in range(candidate_path_len - 1):
                        link_u = candidate_path[path_start_index]
                        link_v = candidate_path[path_start_index + 1]
                        # 将 link_u,link_v Swap
                        if(link_u > link_v):
                            z_ijt -= y[i, j, link_v, link_u, t] / (candidate_path_len - 1)
                        else:
                            z_ijt -= y[i, j, link_u, link_v, t] / (candidate_path_len - 1)
                    model.add_constraint(z_ijt <= 0, ctname="z_less_y")

        # 2d
        for i in range(self.sd_num):
            for t in range(fixed_T):
                w_it = model.linear_expr()
                w_it += w[i, t]
                for j in range(len(self.net.candidate_paths[i]["paths"])):
                    w_it -= z[i, j, t]
                model.add_constraint(w_it <= 0, ctname="w_less_z")

        def isNotInPath(path, edge) -> bool:
            for i in range(len(path) - 1):
                if(edge[0] == path[i] and edge[1] == path[i+1]):
                    return False
                if(edge[1] == path[i] and edge[0] == path[i+1]):
                    return False
            return True

        # 2e
        for i in range(self.sd_num):
            # 改一下 fixed_T 索引
            model.add_constraint(w[i, fixed_T - 1] == 1, ctname="w_equals_1")
        
        # 2f
        for i in range(self.sd_num):
            for j in range(len(self.net.candidate_paths[i]["paths"])):
                candidate_path = self.net.candidate_paths[i]["paths"][j]
                candidate_path_len = len(candidate_path)
                for edge in self.net.G.edges():
                    if(isNotInPath(candidate_path, edge)):
                        if(edge[0]>edge[1]):
                            model.add_constraint(y[i, j, edge[1], edge[0], 0] == 0, ctname="y_equals_0")
                        else:
                            model.add_constraint(y[i, j, edge[0], edge[1], 0] == 0, ctname="y_equals_0_swap")
            
        for i in range(self.sd_num):
            for j in range(len(self.net.candidate_paths[i]["paths"])):
                candidate_path = self.net.candidate_paths[i]["paths"][j]
                candidate_path_len = len(candidate_path)
                for tmp_idx in range(candidate_path_len - 1):
                    link = [candidate_path[tmp_idx], candidate_path[tmp_idx+1]]
                    # 等于 ==> 大于等于
                    # 改成单向的
                    if(link[0]>link[1]):
                        model.add_constraint(y[i, j, link[1], link[0], fixed_T] >= z[i, j, fixed_T], ctname="y_equals_z")
                    else:
                        model.add_constraint(y[i, j, link[0], link[1], fixed_T] >= z[i, j, fixed_T], ctname="y_equals_z")

        objective_T = model.integer_var(lb=0, ub=fixed_T, name="objective_T")

        model.minimize(objective_T)
        # model.export_as_lp("model.lp")
        return model

